use std::{marker::PhantomData, ops::Deref};

use crate::scene::AsBytes;

pub struct DynBuffer<D: AsBytes> {
    buffer: wgpu::Buffer,
    _p: PhantomData<D>,
}

impl<D: AsBytes> DynBuffer<D> {
    pub fn new(ctx: &super::Context, amount: usize, usage: wgpu::BufferUsages) -> Self {
        let buffer = ctx.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: (amount * std::mem::size_of::<D>()) as _,
            mapped_at_creation: false,
            usage: usage | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST,
        });

        Self {
            buffer,
            _p: PhantomData,
        }
    }

    pub fn write(&mut self, ctx: &super::Context, data: &[D]) {
        let bytes = data.as_bytes();
        if bytes.len() as u64 > self.buffer.size() {
            self.resize(ctx, data.len())
        }

        ctx.queue.write_buffer(&self.buffer, 0, bytes)
    }

    pub fn resize(&mut self, ctx: &super::Context, amount: usize) {
        let size = (amount * std::mem::size_of::<D>()) as _;
        if size <= self.buffer.size() {
            return;
        }

        let new_buffer = ctx.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size,
            mapped_at_creation: false,
            usage: self.buffer.usage(),
        });

        if self.buffer.size() > 0 {
            let mut cmd = ctx.device.create_command_encoder(&Default::default());
            cmd.copy_buffer_to_buffer(&self.buffer, 0, &new_buffer, 0, self.buffer.size());
            ctx.queue.submit(std::iter::once(cmd.finish()));
        }

        self.buffer = new_buffer;
    }
}

impl<D: AsBytes> Deref for DynBuffer<D> {
    type Target = wgpu::Buffer;

    fn deref(&self) -> &Self::Target {
        &self.buffer
    }
}

pub struct UniformBuffer<D: AsBytes> {
    buffer: wgpu::Buffer,
    layout: wgpu::BindGroupLayout,
    bind_group: wgpu::BindGroup,
    _p: PhantomData<D>,
}

impl<D: AsBytes> UniformBuffer<D> {
    pub fn new(ctx: &super::Context, visibility: wgpu::ShaderStages) -> Self {
        let layout = ctx
            .device
            .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: None,
                entries: &[wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Uniform,
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                }],
            });

        let buffer = ctx.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: std::mem::size_of::<D>() as u64,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let bind_group = ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &layout,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: buffer.as_entire_binding(),
            }],
        });

        Self {
            layout,
            bind_group,
            buffer,
            _p: PhantomData,
        }
    }

    pub fn write(&mut self, ctx: &super::Context, data: &D) {
        ctx.queue.write_buffer(&self.buffer, 0, data.as_bytes())
    }

    pub fn bind_group(&self) -> &wgpu::BindGroup {
        &self.bind_group
    }

    pub fn layout(&self) -> &wgpu::BindGroupLayout {
        &self.layout
    }
}
